[OMNIML-6076] [13/21] Describe expert parallelism in the model specs - #2746

keep_compressed asked a model's prepare hook to keep its experts packed, for
a frozen speculative-decoding target. PTQ always needs the decompressed
weights and hf_ptq never set it, so the knob and its pass-through go. prepare
now receives ckpt_path and trust_remote_code only, which the test pins.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The applier they exercise enforces each model's transformers floor, which
imports transformers. The sibling model tests skip without it the same way.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

transformers ships a base_model_ep_plan for gpt_oss, so a torchrun run took
the expert-parallel path and failed only at quantization, minutes into the
load. The loader drops GptOssExperts' own quantized wrapper so that the
generic grouped-GEMM one takes over, and that one refuses experts stored
transposed: the fused-experts export cannot write them. It would split
gate and up on the wrong axis, and GPT-OSS also interleaves them.

Add EPSpec.unsupported. apply_distributed_support refuses a model that
names a reason, before anything is applied or loaded, and points at FSDP2.
Mark gpt_oss with it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Per-model expert-parallel support in the same structure as the model specs:
one package per HF model type under modelopt/torch/models, data in specs.py, and
code in ep.py only where the model genuinely needs it. These two need code.

Llama-4 is the one MoE that does not follow the transformers sparse-dispatch
convention -- its block routes inline and computes every expert densely, so
there is no router-to-experts seam for the transport to wrap; ep.py adapts it to
the ExpertsInterface contract instead. llama4_text reuses llama4's adapter and
differs only in where the plan lives -- the config itself rather than a nested
text config, since there is no VLM wrapper.

MiniMax-M3 already follows the convention, with fused [E, out, in] expert
params, but hardcodes a per-expert loop in its experts forward, so it only needs
its experts class routed through the ExpertsInterface.

The declarative half is covered by the spec tests. The prepare hooks patch real
transformers classes, so they are exercised by the PTQ runs rather than by unit
tests.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The Llama-4 adapter rewrote Llama4TextMoe.forward to dispatch its top-k
through the experts interface, but it scaled each expert's output by the
router score where Llama-4 scales the input, and SwiGLU does not commute
with that: on a tiny config the outputs differ by up to 6.9 on values near
11. Its experts are also stored transposed, which the expert-parallel export
cannot write, so the run failed at quantization anyway. Drop the adapter and
refuse llama4 and llama4_text the way gpt_oss is refused: transformers ships
a plan for both, which would otherwise route them into the expert-parallel
path.

Every transformers release with minimax_m3_vl already decorates
MiniMaxM3VLExperts exactly as the M3 adapter did and ships the text
config's base_model_ep_plan, and the loader selects the all-to-all experts
forward after loading, so the adapter changed nothing at runtime. Keep only
its transformers floor.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

nemotron_h/ep.py registered a NemotronHExpertsEP parallel style for the
plan's nemotron_experts_ep entry, but load_and_prepare_ep maps that entry to
torch_a2a_experts before loading, so the style never ran. Registering it
imported transformers' all_reduce_forward, which 5.18 no longer has, so
applying Nemotron-H's spec failed on CI's latest transformers; its _prepare_*
hooks are not called from 5.16 on either.

The plan moves into the spec that both spellings of the architecture share,
with no prepare hook.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

transformers 5.18 added nemotron_h_omni, so its spec can no longer claim remote code;
test_specs_vs_transformers fails on the latest transformers until it names the release.
Expert parallelism still needs the checkpoint's code, whose llm_config the plan targets.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

For several models the spec is what supplies the expert-parallel plan: DeepSeek-V3 ships none of
its own, so deciding without the spec would read its MoE as dense. Unless the run rules expert
parallelism out (a single process, or an EP degree of 1), the model's spec is applied to the config
before rank 0 builds the model to decide, and the returned config carries it. A load prepares that
config once, and not at all when handed the one resolve_distribution_plan returned, since a model's
preparation can patch its modeling code. A run that rules expert parallelism out applies no spec,
so a model whose spec refuses expert parallelism, like GPT-OSS, still loads FSDP2-sharded.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Each rank read its own MODELOPT_EP_DEGREE to decide whether to apply
the model spec, though the knobs are rank 0's. A rank whose environment
ruled expert parallelism out skipped the spec rank 0 applied, loaded
without the expert-parallel plan the others shard by, and raised once
the mesh was up, leaving them waiting. Rank 0 decides it and sends it,
as it does the plan.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

An EP degree of 1 skips the model's spec, so the guard against leaving an
expert-parallel MoE whole on every rank found routed experts only through
the config's own base_model_ep_plan. DeepSeek-V3, whose plan comes only
from its spec, therefore loaded whole on every rank at MODELOPT_EP_DEGREE=1
with a replicated dense part, where Qwen3-MoE, which carries its plan,
refused. The guard now also looks for the experts through the plan the
spec would supply. The spec itself stays skipped: it refuses models that
cannot be expert-parallel, which an EP degree of 1 is how to load.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

At an EP degree of 1 the guard against holding an expert-parallel MoE
whole on every rank falls back on the plan the model's spec would supply,
since that degree skips the spec. It read only the spec's
base_model_ep_plan, so a remote-code model whose prepare attaches the
plan, after patching the checkpoint's modeling, showed no routed experts
and loaded whole on every rank. Running prepare there is not an option:
it patches the model's code, for a run that is not expert-parallel.

EPSpec.prepared_ep_plan states the plan prepare attaches, and the plan
the guard reads includes it. prepare still attaches the plan itself; the
spec only states it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The refusal of an MoE held whole on every rank told every model to raise
MODELOPT_EP_DEGREE above 1. GPT-OSS and Llama-4 carry plans of their own,
so they reach that refusal too, but their specs refuse expert
parallelism: following the advice would only have hit that refusal next.
For them the error now says they cannot be quantized expert-parallel yet
and points at FSDP2 alone, as their specs' own refusal does. Both
messages name the ways to ask for FSDP2: --use_fsdp2 in hf_ptq, or
MODELOPT_FSDP_SIZE=world.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…ided

The test played rank 1 against a stand-in broadcast that answered with
rank 0's decision. The pooled job that checks every rank takes rank 0's
plan now covers the spec too: rank 1's own EP knob says the opposite of
rank 0's, and rank 1 must apply DeepSeek-V3's spec exactly when rank 0
did, over the real broadcast, or it would load without the plan the
others shard by.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

test_the_plan_a_spec_would_supply imported modelopt.torch.models and
modelopt.torch.utils.moe inside the test, for no circular-import or
optional-dependency reason: the module's transformers guard already
precedes its other modelopt imports.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

- test_some_model_declares_expert_parallel_support goes: with no spec
  registered, the named-model applier and decision tests fail first.
- test_a_model_type_with_no_spec_is_left_alone goes: every multi-rank
  dense decision test runs the same early return.
- test_the_plan_is_attached_when_the_config_has_none goes:
  test_the_spec_is_applied_before_the_decision attaches the real
  deepseek_v3 plan.
- The refusal test is parametrized over gpt_oss, llama4 and llama4_text by
  name rather than over whichever specs refuse, so a lost refusal fails it
  instead of shrinking it. It replaces
  test_models_whose_own_plan_would_misroute_them_are_refused.
- The two prepare probes are one test.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

modelopt.torch.models exports EPSpec, but specs.py, where it is defined,
left it out of its __all__.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The expert-parallel route of a sharded load handed a plan but no config
read the config without token, cache_dir or local_files_only, as the
FSDP2 route did before. It now reads it with the Hub options the load was
given.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

A config passed for the load is the caller's own. With the plan
decided by the load, resolve_distribution_plan prepares a copy of it
through the model spec. With a plan handed in, the expert-parallel load
took the config raw, so the spec's preparation never ran: DeepSeek-V3
failed for want of the expert-parallel plan the spec adds, and a model
the spec refuses was not refused. The load now prepares a copy, once,
as deciding does.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The config resolve_distribution_plan prepared comes back to the load
with its plan, and the load does not prepare it again: preparing can
patch modeling code. It skipped the spec's checks with the preparation,
so a hand-made plan and config loaded GPT-OSS or Llama-4 expert-parallel
although their specs refuse it. check_distributed_support now runs the
spec's refusal and transformers floor on their own, and
apply_distributed_support starts with it.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…sformers floor

A multi-rank run that leaves the EP degree to auto checks the model's spec before its plan is
decided, so a model whose expert parallelism needs a newer transformers (Qwen3-MoE below 5.7)
stops there. Like the refusal of a model that cannot be expert-parallel yet, the error now says it
can also be quantized without expert parallelism.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>